"""
安全工具模块
提供账户锁定、频率限制等安全功能
"""
import re
from datetime import datetime, timedelta
from typing import Optional, Dict, Any
from motor.motor_asyncio import AsyncIOMotorDatabase
import redis.asyncio as redis
import logging

logger = logging.getLogger(__name__)

class SecurityValidator:
    """安全验证器"""
    
    # 账户锁定配置
    MAX_LOGIN_ATTEMPTS = 5  # 最大登录尝试次数
    LOCKOUT_DURATION = 30  # 锁定时长（分钟）
    
    # 频率限制配置
    RATE_LIMIT_WINDOW = 15  # 频率限制窗口（分钟）
    MAX_REQUESTS_PER_WINDOW = 10  # 窗口内最大请求次数
    
    @staticmethod
    def validate_input_format(username: str, password: str) -> Dict[str, Any]:
        """
        验证输入格式
        """
        errors = []
        
        # 检查空值
        if not username or not username.strip():
            errors.append("用户名不能为空")
        
        if not password or not password.strip():
            errors.append("密码不能为空")
            
        # 检查用户名格式（邮箱或用户名）
        if username and username.strip():
            username = username.strip()
            # 检查是否为邮箱格式
            email_pattern = r'^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}$'
            # 检查是否为用户名格式（字母数字下划线，3-20位）
            username_pattern = r'^[a-zA-Z0-9_]{3,20}$'
            
            if not (re.match(email_pattern, username) or re.match(username_pattern, username)):
                errors.append("用户名格式不正确，请输入有效的用户名或邮箱")
        
        # 检查密码长度
        if password and len(password.strip()) < 6:
            errors.append("密码长度不能少于6位")
            
        return {
            "valid": len(errors) == 0,
            "errors": errors,
            "sanitized_username": username.strip() if username else "",
            "sanitized_password": password.strip() if password else ""
        }
    
    @staticmethod
    async def check_account_lockout(
        db: AsyncIOMotorDatabase,
        user_id: str,
        ip_address: Optional[str] = None
    ) -> Dict[str, Any]:
        """
        检查账户是否被锁定
        """
        try:
            # 检查用户级别的锁定
            user_lockout = await db.user_security.find_one({
                "user_id": user_id,
                "type": "account_lockout"
            })
            
            if user_lockout:
                lockout_until = user_lockout.get("locked_until")
                if lockout_until and datetime.utcnow() < lockout_until:
                    remaining_time = lockout_until - datetime.utcnow()
                    return {
                        "locked": True,
                        "reason": "account",
                        "remaining_minutes": int(remaining_time.total_seconds() / 60),
                        "message": f"账户已被锁定，请在 {int(remaining_time.total_seconds() / 60)} 分钟后重试"
                    }
                else:
                    # 锁定已过期，清除记录
                    await db.user_security.delete_one({
                        "user_id": user_id,
                        "type": "account_lockout"
                    })
            
            # 检查IP级别的锁定
            if ip_address:
                ip_lockout = await db.user_security.find_one({
                    "ip_address": ip_address,
                    "type": "ip_lockout"
                })
                
                if ip_lockout:
                    lockout_until = ip_lockout.get("locked_until")
                    if lockout_until and datetime.utcnow() < lockout_until:
                        remaining_time = lockout_until - datetime.utcnow()
                        return {
                            "locked": True,
                            "reason": "ip",
                            "remaining_minutes": int(remaining_time.total_seconds() / 60),
                            "message": f"IP地址已被锁定，请在 {int(remaining_time.total_seconds() / 60)} 分钟后重试"
                        }
                    else:
                        # 锁定已过期，清除记录
                        await db.user_security.delete_one({
                            "ip_address": ip_address,
                            "type": "ip_lockout"
                        })
            
            return {"locked": False}
            
        except Exception as e:
            logger.error(f"检查账户锁定状态时出错: {e}")
            return {"locked": False}
    
    @staticmethod
    async def record_login_attempt(
        db: AsyncIOMotorDatabase,
        user_id: str,
        ip_address: Optional[str],
        success: bool,
        username: str = ""
    ) -> None:
        """
        记录登录尝试
        """
        try:
            current_time = datetime.utcnow()
            
            # 记录登录尝试
            attempt_record = {
                "user_id": user_id,
                "ip_address": ip_address,
                "username": username,
                "success": success,
                "timestamp": current_time,
                "type": "login_attempt"
            }
            
            await db.user_security.insert_one(attempt_record)
            
            # 如果登录失败，检查是否需要锁定账户
            if not success:
                await SecurityValidator._check_and_apply_lockout(
                    db, user_id, ip_address, current_time
                )
            else:
                # 登录成功，清除失败记录
                await SecurityValidator._clear_failed_attempts(db, user_id, ip_address)
                
        except Exception as e:
            logger.error(f"记录登录尝试时出错: {e}")
    
    @staticmethod
    async def _check_and_apply_lockout(
        db: AsyncIOMotorDatabase,
        user_id: str,
        ip_address: Optional[str],
        current_time: datetime
    ) -> None:
        """
        检查并应用账户锁定
        """
        try:
            # 检查最近的失败尝试次数
            time_window = current_time - timedelta(minutes=SecurityValidator.LOCKOUT_DURATION)
            
            # 用户级别的失败次数
            user_failures = await db.user_security.count_documents({
                "user_id": user_id,
                "success": False,
                "timestamp": {"$gte": time_window},
                "type": "login_attempt"
            })
            
            if user_failures >= SecurityValidator.MAX_LOGIN_ATTEMPTS:
                # 锁定账户
                lockout_until = current_time + timedelta(minutes=SecurityValidator.LOCKOUT_DURATION)
                await db.user_security.update_one(
                    {"user_id": user_id, "type": "account_lockout"},
                    {
                        "$set": {
                            "user_id": user_id,
                            "type": "account_lockout",
                            "locked_until": lockout_until,
                            "locked_at": current_time,
                            "reason": f"连续 {user_failures} 次登录失败"
                        }
                    },
                    upsert=True
                )
                logger.warning(f"账户 {user_id} 因连续失败登录被锁定至 {lockout_until}")
            
            # IP级别的失败次数
            if ip_address:
                ip_failures = await db.user_security.count_documents({
                    "ip_address": ip_address,
                    "success": False,
                    "timestamp": {"$gte": time_window},
                    "type": "login_attempt"
                })
                
                if ip_failures >= SecurityValidator.MAX_LOGIN_ATTEMPTS * 2:  # IP锁定阈值更高
                    # 锁定IP
                    lockout_until = current_time + timedelta(minutes=SecurityValidator.LOCKOUT_DURATION)
                    await db.user_security.update_one(
                        {"ip_address": ip_address, "type": "ip_lockout"},
                        {
                            "$set": {
                                "ip_address": ip_address,
                                "type": "ip_lockout",
                                "locked_until": lockout_until,
                                "locked_at": current_time,
                                "reason": f"IP地址连续 {ip_failures} 次登录失败"
                            }
                        },
                        upsert=True
                    )
                    logger.warning(f"IP地址 {ip_address} 因连续失败登录被锁定至 {lockout_until}")
                    
        except Exception as e:
            logger.error(f"检查和应用账户锁定时出错: {e}")
    
    @staticmethod
    async def _clear_failed_attempts(
        db: AsyncIOMotorDatabase,
        user_id: str,
        ip_address: Optional[str]
    ) -> None:
        """
        清除失败尝试记录
        """
        try:
            # 清除用户的锁定状态
            await db.user_security.delete_many({
                "user_id": user_id,
                "type": {"$in": ["account_lockout"]}
            })
            
            # 清除IP的锁定状态（如果该IP只有这个用户的失败记录）
            if ip_address:
                await db.user_security.delete_many({
                    "ip_address": ip_address,
                    "type": {"$in": ["ip_lockout"]}
                })
                
        except Exception as e:
            logger.error(f"清除失败尝试记录时出错: {e}")
            # 清除用户的锁定状态
            await db.user_security.delete_many({
                "user_id": user_id,
                "type": {"$in": ["account_lockout"]}
            })
            
            # 清除IP的锁定状态（如果该IP只有这个用户的失败记录）
            if ip_address:
                await db.user_security.delete_many({
                    "ip_address": ip_address,
                    "type": {"$in": ["ip_lockout"]}
                })
                
        except Exception as e:
            logger.error(f"清除失败尝试记录时出错: {e}")
    
    @staticmethod
    async def check_rate_limit(
        redis_client: Optional[redis.Redis],
        ip_address: Optional[str],
        user_identifier: str = ""
    ) -> Dict[str, Any]:
        """
        检查频率限制
        """
        if not redis_client or not ip_address:
            return {"limited": False}
        
        try:
            current_time = datetime.utcnow()
            window_start = current_time - timedelta(minutes=SecurityValidator.RATE_LIMIT_WINDOW)
            
            # 使用IP地址作为主要标识符
            rate_limit_key = f"rate_limit:{ip_address}"
            
            # 获取当前窗口内的请求次数
            request_count = await redis_client.get(rate_limit_key)
            request_count = int(request_count) if request_count else 0
            
            if request_count >= SecurityValidator.MAX_REQUESTS_PER_WINDOW:
                return {
                    "limited": True,
                    "message": f"请求过于频繁，请在 {SecurityValidator.RATE_LIMIT_WINDOW} 分钟后重试",
                    "retry_after": SecurityValidator.RATE_LIMIT_WINDOW * 60
                }
            
            # 增加请求计数
            await redis_client.incr(rate_limit_key)
            await redis_client.expire(rate_limit_key, SecurityValidator.RATE_LIMIT_WINDOW * 60)
            
            return {"limited": False}
            
        except Exception as e:
            logger.error(f"检查频率限制时出错: {e}")
            return {"limited": False}
    
    @staticmethod
    async def cleanup_expired_records(db: AsyncIOMotorDatabase) -> None:
        """
        清理过期的安全记录
        """
        try:
            current_time = datetime.utcnow()
            
            # 清理过期的登录尝试记录（保留30天）
            expire_time = current_time - timedelta(days=30)
            await db.user_security.delete_many({
                "type": "login_attempt",
                "timestamp": {"$lt": expire_time}
            })
            
            # 清理过期的锁定记录
            await db.user_security.delete_many({
                "type": {"$in": ["account_lockout", "ip_lockout"]},
                "locked_until": {"$lt": current_time}
            })
            
            logger.info("安全记录清理完成")
            
        except Exception as e:
            logger.error(f"清理过期安全记录时出错: {e}")